package context

import (
	"time"

	"github.com/yaoapp/gou/connector/openai"
	"github.com/yaoapp/yao/agent/output"
	"github.com/yaoapp/yao/agent/output/message"
)

// Send sends a message via the output module
// Automatically manages BlockID, ThreadID, lifecycle events, and metadata for delta operations
// - For delta operations: inherits BlockID and ThreadID from original message, increments chunk count
// - For new messages: auto-sets ThreadID from Stack, sends message_start event
// - Sends block_start event when a new BlockID is first encountered
// - Records metadata for all sent messages to enable delta inheritance
func (ctx *Context) Send(msg *message.Message) error {
	out, err := ctx.getOutput()
	if err != nil {
		return err
	}

	// Skip lifecycle events for event-type messages (prevent recursion)
	isEventMessage := msg.Type == message.TypeEvent

	// === Handle message_start event: record metadata for future delta chunks ===
	if isEventMessage && msg.Props != nil {
		if event, ok := msg.Props["event"].(string); ok && event == message.EventMessageStart {
			if data, ok := msg.Props["data"].(message.EventMessageStartData); ok {
				// Record metadata from message_start event
				if data.MessageID != "" && ctx.messageMetadata != nil {
					ctx.messageMetadata.setMessage(data.MessageID, &MessageMetadata{
						MessageID:  data.MessageID,
						ThreadID:   data.ThreadID,
						Type:       data.Type,
						StartTime:  time.Now(),
						ChunkCount: 0, // Will be incremented by delta chunks
					})
				}
			}
		}
	}

	// === Delta operations: Auto-inherit and update metadata ===
	if msg.Delta && msg.MessageID != "" && ctx.messageMetadata != nil {
		if metadata := ctx.getMessageMetadata(msg.MessageID); metadata != nil {
			// Inherit BlockID if not specified
			if msg.BlockID == "" {
				msg.BlockID = metadata.BlockID
			}
			// Inherit ThreadID if not specified
			if msg.ThreadID == "" {
				msg.ThreadID = metadata.ThreadID
			}

			// Increment chunk count for this message
			metadata.ChunkCount++
		}
	}

	// === Auto-generate ChunkID (always) ===
	if msg.ChunkID == "" && !isEventMessage {
		if ctx.IDGenerator != nil {
			msg.ChunkID = ctx.IDGenerator.GenerateChunkID()
		} else {
			msg.ChunkID = message.GenerateNanoID()
		}
	}

	// === Non-delta operations: New message logic ===
	if !msg.Delta && !isEventMessage {
		// Auto-set ThreadID for non-root Stack (nested agent calls)
		if msg.ThreadID == "" && ctx.Stack != nil && !ctx.Stack.IsRoot() {
			msg.ThreadID = ctx.Stack.ID
		}

		// BlockID is NOT auto-generated by default (only manually specified in special cases)
		// Example: Send a web card after LLM output, group them in the same Block
		// Developers can specify via ctx.Send(message, blockId) or message.block_id

		// === Send block_start event if this is a new block ===
		if msg.BlockID != "" && ctx.messageMetadata != nil {
			if ctx.messageMetadata.getBlock(msg.BlockID) == nil {
				// New block, send block_start event
				blockStartData := message.EventBlockStartData{
					BlockID:   msg.BlockID,
					Type:      "mixed", // Default type, can be enhanced later
					Timestamp: time.Now().UnixMilli(),
				}
				blockStartEvent := output.NewEventMessage(message.EventBlockStart, "Block started", blockStartData)
				if err := ctx.sendRaw(blockStartEvent); err != nil {
					return err
				}

				// Record block metadata
				ctx.messageMetadata.setBlock(msg.BlockID, &BlockMetadata{
					BlockID:      msg.BlockID,
					Type:         "mixed",
					StartTime:    time.Now(),
					MessageCount: 0,
				})
			}

			// Increment message count for this block
			ctx.messageMetadata.updateBlock(msg.BlockID, func(block *BlockMetadata) {
				block.MessageCount++
			})
		}

		// === Generate MessageID if not provided ===
		if msg.MessageID == "" {
			if ctx.IDGenerator != nil {
				msg.MessageID = ctx.IDGenerator.GenerateMessageID()
			} else {
				msg.MessageID = message.GenerateNanoID() // Use NanoID generator
			}
		}

		// === Send message_start event ===
		messageStartData := message.EventMessageStartData{
			MessageID: msg.MessageID,
			Type:      msg.Type,
			Timestamp: time.Now().UnixMilli(),
			ThreadID:  msg.ThreadID, // Include ThreadID for concurrent stream identification
		}
		messageStartEvent := output.NewEventMessage(message.EventMessageStart, "Message started", messageStartData)
		if err := ctx.sendRaw(messageStartEvent); err != nil {
			return err
		}

		// === Record message metadata with start time ===
		if ctx.messageMetadata != nil {
			ctx.messageMetadata.setMessage(msg.MessageID, &MessageMetadata{
				MessageID:  msg.MessageID,
				BlockID:    msg.BlockID,
				ThreadID:   msg.ThreadID,
				Type:       msg.Type,
				StartTime:  time.Now(),
				ChunkCount: 1, // Initial chunk
			})
		}
	}

	// === Actually send the message ===
	if err := out.Send(msg); err != nil {
		return err
	}

	// === Auto-send message_end for non-delta messages (complete messages) ===
	if !msg.Delta && !isEventMessage && msg.MessageID != "" && ctx.messageMetadata != nil {
		metadata := ctx.messageMetadata.getMessage(msg.MessageID)
		if metadata != nil {
			// Calculate duration
			durationMs := time.Since(metadata.StartTime).Milliseconds()

			// Extract content for the extra field
			var content interface{}
			if msg.Props != nil {
				if c, ok := msg.Props["content"]; ok {
					content = c
				}
			}

			// Build message_end event data
			endData := message.EventMessageEndData{
				MessageID:  msg.MessageID,
				Type:       msg.Type,
				Timestamp:  time.Now().UnixMilli(),
				ThreadID:   metadata.ThreadID, // Include ThreadID for concurrent stream identification
				DurationMs: durationMs,
				ChunkCount: metadata.ChunkCount,
				Status:     "completed",
			}

			// Add content to extra if available
			if content != nil {
				endData.Extra = map[string]interface{}{
					"content": content,
				}
			}

			// Send message_end event
			messageEndEvent := output.NewEventMessage(message.EventMessageEnd, "Message completed", endData)
			ctx.sendRaw(messageEndEvent)
		}
	}

	return nil
}

// EndMessage sends a message_end event for a completed message
// Note: For non-delta messages, message_end is automatically sent by Send()
// This method is primarily for delta streaming scenarios:
// - After all delta chunks are sent for a message, call EndMessage() to finalize it
// - For LLM streaming, this is typically called after receiving ChunkMessageEnd
func (ctx *Context) EndMessage(messageID string, content interface{}) error {
	if messageID == "" || ctx.messageMetadata == nil {
		return nil
	}

	metadata := ctx.messageMetadata.getMessage(messageID)
	if metadata == nil {
		return nil // Message not found, skip
	}

	// Calculate duration
	durationMs := time.Since(metadata.StartTime).Milliseconds()

	// Build message_end event data
	endData := message.EventMessageEndData{
		MessageID:  messageID,
		Type:       metadata.Type,
		Timestamp:  time.Now().UnixMilli(),
		ThreadID:   metadata.ThreadID, // Include ThreadID for concurrent stream identification
		DurationMs: durationMs,
		ChunkCount: metadata.ChunkCount,
		Status:     "completed",
	}

	// Add content to extra if provided
	if content != nil {
		endData.Extra = map[string]interface{}{
			"content": content,
		}
	}

	// Send message_end event
	messageEndEvent := output.NewEventMessage(message.EventMessageEnd, "Message completed", endData)
	return ctx.sendRaw(messageEndEvent)
}

// EndBlock sends a block_end event for a completed block
// This should be called explicitly when all messages in a block are complete
func (ctx *Context) EndBlock(blockID string) error {
	if blockID == "" || ctx.messageMetadata == nil {
		return nil
	}

	blockMetadata := ctx.messageMetadata.getBlock(blockID)
	if blockMetadata == nil {
		return nil // Block not found, skip
	}

	// Calculate duration
	durationMs := time.Since(blockMetadata.StartTime).Milliseconds()

	// Build block_end event data
	endData := message.EventBlockEndData{
		BlockID:      blockID,
		Type:         blockMetadata.Type,
		Timestamp:    time.Now().UnixMilli(),
		DurationMs:   durationMs,
		MessageCount: blockMetadata.MessageCount,
		Status:       "completed",
	}

	// Send block_end event
	blockEndEvent := output.NewEventMessage(message.EventBlockEnd, "Block completed", endData)
	return ctx.sendRaw(blockEndEvent)
}

// SendGroup sends a group of messages via the output module
// Deprecated: This method is deprecated and will be removed in future versions
func (ctx *Context) SendGroup(group *message.Group) error {
	output, err := ctx.getOutput()
	if err != nil {
		return err
	}
	return output.SendGroup(group)
}

// Flush flushes the output writer
func (ctx *Context) Flush() error {
	output, err := ctx.getOutput()
	if err != nil {
		return err
	}
	return output.Flush()
}

// CloseOutput closes the output writer
func (ctx *Context) CloseOutput() error {
	output, err := ctx.getOutput()
	if err != nil {
		return err
	}
	return output.Close()
}

// sendRaw sends a message directly without triggering lifecycle events
// Used internally to send event messages without recursion
func (ctx *Context) sendRaw(msg *message.Message) error {
	out, err := ctx.getOutput()
	if err != nil {
		return err
	}
	return out.Send(msg)
}

// getWriter gets the effective Writer for the current context
// Priority: Skip.Output > Stack.Options.Writer > ctx.Writer
func (ctx *Context) getWriter() Writer {
	// Check if output is explicitly skipped (for internal A2A calls)
	if ctx.Stack != nil && ctx.Stack.Options != nil && ctx.Stack.Options.Skip != nil && ctx.Stack.Options.Skip.Output {
		return nil // Explicitly disable output
	}

	// Check if current Stack has a Writer override
	if ctx.Stack != nil && ctx.Stack.Options != nil && ctx.Stack.Options.Writer != nil {
		return ctx.Stack.Options.Writer
	}

	return ctx.Writer
}

// getOutput gets the output writer for the context
func (ctx *Context) getOutput() (*output.Output, error) {
	// Check if current Stack has cached output
	if ctx.Stack != nil && ctx.Stack.output != nil {
		return ctx.Stack.output, nil
	}

	trace, _ := ctx.Trace()
	var options message.Options = message.Options{
		BaseURL: "/",
		Writer:  ctx.getWriter(), // Use getWriter() to resolve Writer priority
		Trace:   trace,
		Locale:  ctx.Locale,
		Accept:  string(ctx.Accept),
	}

	// Set ModelCapabilities (now using openai.Capabilities directly)
	if ctx.Capabilities != nil {
		caps := openai.Capabilities(*ctx.Capabilities)
		options.Capabilities = &caps
	}

	out, err := output.NewOutput(options)
	if err != nil {
		return nil, err
	}

	// Cache to current Stack (each Stack has its own output with its own Writer)
	if ctx.Stack != nil {
		ctx.Stack.output = out
	}

	return out, nil
}
